#!/usr/bin/python

import sys
import numpy as np
import scipy

def ConvertPLY(fout, depth, color = None):
    '''
    This function write depth and color out to a ply format file.
    Assume color is a 3-D array with last element represents RGB
    (R = 0, G = 1, B = 2).
    This function will ignore all points that has a negtive depth.
    '''
    assert depth.ndim == 2
    if __debug__:
        if color != None:
            if not color.ndim == 3:
                raise AssertionError
            (x, y, z) = color.shape
            if not z == 3:
                raise AssertionError
            if not depth.shape == (x, y):
                raise AssertionError

    # get dimensions
    (rows, cols) = depth.shape
    
    # if color not given, assume grey (0x808080)
    if color == None:
        color = np.empty((rows, cols, 3))
        color[:,:,:] = 128
    
    # vertex index table
    vertex = np.empty((rows, cols))
    # output line buffer
    output = []
    
    # prepare header
    output.append("ply\n")
    output.append("format ascii 1.0\n")
    output.append("") # 2
    output.append("property float x\n")
    output.append("property float y\n")
    output.append("property float z\n")
    output.append("property uchar red\n")
    output.append("property uchar green\n")
    output.append("property uchar blue\n")
    output.append("") # 9
    output.append("property list uchar int vertex_index\n")
    output.append("end_header\n")
    
    # prepare vertex info
    vertexindex = 0
    for row in range(rows):
        for col in range(cols):
            if depth[row, col] >= 0:
                vertex[row, col] = vertexindex
                vertexindex += 1
                output.append("%f %f %f %d %d %d\n" % \
                           (col, row, depth[row, col], \
                            color[row, col, 0], \
                            color[row, col, 1], \
                            color[row, col, 2]))
            else:
                vertex[row, col] = -1
    
    # prepare face info
    faceCount = 0
    for row in range(rows-1):
        for col in range(cols-1):
            # a b
            # c d
            # triangle acb & bcd
            a = vertex[row, col]
            b = vertex[row, col+1]
            c = vertex[row+1, col]
            d = vertex[row+1, col+1]
            if a < 0:
                if b >= 0 and c >= 0 and d >= 0:
                    output.append("3 %d %d %d\n" %(b, c, d))
                    faceCount += 1
            elif b < 0: # a >= 0
                if c >= 0 and d >= 0:
                    output.append("3 %d %d %d\n" %(a, c, d))
                    faceCount += 1
            elif c < 0: # a, b >= 0
                if d >= 0:
                    output.append("3 %d %d %d\n" %(a, d, b))
                    faceCount += 1
            elif d < 0: # a, b, c >= 0
                output.append("3 %d %d %d\n" %(a, c, b))
                faceCount += 1
            else:       # a, b, c, d >= 0
                output.append("3 %d %d %d\n" %(a, c, b))
                output.append("3 %d %d %d\n" %(b, c, d))
                faceCount += 2
    
    # update element info
    output[2] = "element vertex %d\n" % vertexindex
    output[9] = "element face %d\n" % faceCount
    
    # write output to file
    fout.writelines(output)

if __name__ == "__main__":
    argc = len(sys.argv)
    if argc < 3 or argc > 4:
        print "Usage: ConvertPLY outputfile depthfile [colorfile]"
        sys.exit(1)
    
    fout = open(sys.argv[1], 'w')
    if argc == 3:
        depth = scipy.loadtxt(sys.argv[2])
        ConvertPLY(fout, depth)
    else: # argc == 4
        depth = scipy.loadtxt(sys.argv[2])
        color = scipy.loadtxt(sys.argv[3])
        ConvertPLY(fout, color)
    fout.close()

